package org.pz.pojo;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Mono;

import java.util.Date;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 令牌桶限流算法，来自 deepSeek
 * 现在只能是单体的 gateway 限流，想要分布式限流，需要吧参数存在 redis 一类的限流中
 * */
@Slf4j
public class TokenBucketRateLimiter {

    private final long capacity;          // 桶容量（最大令牌数）
    private final long refillTokens;      // 每次填充的令牌数
    private final long refillPeriodMillis;// 填充周期（毫秒）

    // 原子变量记录当前令牌数和上次填充时间
    private final AtomicLong tokens;
    private final AtomicReference<Long> lastRefillTime;

    public TokenBucketRateLimiter(long capacity, long refillTokens, long refillPeriodMillis) {
        this.capacity = capacity;
        this.refillTokens = refillTokens;
        this.refillPeriodMillis = refillPeriodMillis;
        this.tokens = new AtomicLong(capacity);  // 初始化令牌桶
        this.lastRefillTime = new AtomicReference<>(System.currentTimeMillis());
    }

    /**
     * 尝试获取一个令牌
     * @return true-允许通过；false-限流
     */
    public Mono<Boolean> tryAcquire() {
        log.info("当前令牌数：{}, 上次填充时间：{}", tokens, new Date(lastRefillTime.get()));
        return Mono.fromSupplier(() -> {
            long now = System.currentTimeMillis();
            Long oldLastTime = lastRefillTime.get();

            // 计算时间差，判断是否需要填充令牌
            long timeElapsed = now - oldLastTime;
            if (timeElapsed > refillPeriodMillis) {
                // CAS 更新上次填充时间，避免并发重复计算
                if (lastRefillTime.compareAndSet(oldLastTime, now)) {
                    long newTokens = Math.min(capacity, tokens.get() + refillTokens);
                    tokens.set(newTokens);
                }
            }

            // 尝试消耗令牌（非阻塞方式）
            long currentTokens = tokens.get();
            if (currentTokens > 0) {
                return tokens.compareAndSet(currentTokens, currentTokens - 1);
            }
            return false;
        });
    }
}
